//============================================================================
// Copyright (C) 2010-2013 Brett R. Jones
//
// Issued to public domain 2010
//
// brettjones1900@gmail.com
// http://www.p2panarchy.com
//============================================================================
#include "VxUtilLib.h"
#include "CheezyCrypt.h"
#include <memory.h>
#include <string>
#ifdef _MSC_VER
	#include <WinSock2.h>
	#include "windows.h"
#else
	#include <sys/time.h>
	#include <cstdlib>
#endif // _MSC_VER
#include "md5.h"

//=== cheezy crypto with known weeknesses ===//
// issued to public domain year 2006
//===========================================//

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//=== constructor ===//
CheezyKey::CheezyKey()
	: m_bIsSet(0)
{
}

//---------------------------------------------------------------------------
//=== destructor ===//
CheezyKey::~CheezyKey()
{
	memset( m_au32Key, 0xa5, sizeof( m_au32Key ) );
}

//---------------------------------------------------------------------------
//! set key from data..
RCODE CheezyKey::ImportKey( unsigned char * pu8Data, int iLen )
{
	//vx_assert( iLen == sizeof( m_au8Key ) );
	if( iLen == sizeof( m_au32Key ) )
	{
		memcpy( m_au32Key, pu8Data, iLen );
		m_bIsSet = true;
		return 0;
	}
	return -1;
}
//---------------------------------------------------------------------------
//! set encryption key from another
void CheezyKey::ImportKey( VxKey * poKey )
{ 
	memcpy( m_au32Key, poKey->m_au32Key, sizeof(m_au32Key ) ); 
	m_bIsSet = poKey->m_bIsSet; 
}

//---------------------------------------------------------------------------
//! return key into data buffer
RCODE CheezyKey::ExportKey( unsigned char * pu8RetKeyData, int iBufLen )
{
	if( iBufLen >= (int)sizeof( m_au32Key ) )
	{
		memcpy( pu8RetKeyData, m_au32Key, sizeof( m_au32Key ) );
		return 0;
	}
	return -1;
}
//---------------------------------------------------------------------------
//! export encryption key to another another
void CheezyKey::ExportKey( VxKey * poKey )
{ 
	memcpy( poKey->m_au32Key, m_au32Key, sizeof(m_au32Key ) ); 
	poKey->m_bIsSet = m_bIsSet ; 
}

//---------------------------------------------------------------------------
//! make encryption key from user name and password
RCODE CheezyKey::SetKeyFromPassword(	const char *	pUserName,			// user name
										const char *	pPassword,			// password
										const char *	pSalt )				// salt
{
	std::string strPwd = pUserName;
	strPwd += pPassword;

	return SetKeyFromPassword(	strPwd.c_str(),	(int)strPwd.size(), pSalt );		
}


//---------------------------------------------------------------------------
//! make encryption key from password
RCODE CheezyKey::SetKeyFromPassword(	const char *	pPassword,			// password
										int				iPasswordLen,		// length of password	int iPasswordLen )
										const char *	pSalt )	// salt
{
	//vx_assert( pPassword );
	//vx_assert( (iPasswordLen > 0) && (iPasswordLen < 256) );
    struct MD5Context   md5c;

	if((pPassword == NULL) || ( 0 == iPasswordLen ) )
    {
        return -1;
    }
    MD5Init(&md5c);
    MD5Update(&md5c, (unsigned char *)pPassword, (unsigned int)iPasswordLen);
    MD5Final((unsigned char *)m_au32Key, &md5c);
	m_bIsSet = true;
	return 0;
}
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
//=== constructor ===//
CheezyCrypt::CheezyCrypt()
	: m_bIsKeyValid(0)
{ 
	// empty
}


//---------------------------------------------------------------------------
//! generate key from password and set encryption key in one fuction call
//! NOTE: Max password len 255
RCODE CheezyCrypt::SetPassword( const char * pPassword, int iPasswordLen )
{
    struct MD5Context   md5c;
    unsigned char       bfvec[CHEEZY_SYM_KEY_LEN];

    if(pPassword == NULL)
    {
        return -1;
    }
    MD5Init(&md5c);
    MD5Update(&md5c, (unsigned char *)pPassword, (unsigned int)iPasswordLen);
    MD5Final(bfvec, &md5c);
	BlowSetKey( &m_oBlowCtx, bfvec, CHEEZY_SYM_KEY_LEN );
	m_bIsKeyValid = true;
	return 0;
}
//---------------------------------------------------------------------------
//! Generate encryption key from password
RCODE CheezyCrypt::GenerateKey( const char * pPassword, int iPasswordLen, CheezyKey * poRetKey )
{
    struct MD5Context   md5c;
    unsigned char       bfvec[CHEEZY_SYM_KEY_LEN];

    if(pPassword == NULL)
    {
        return -1;
    }
    MD5Init(&md5c);
    MD5Update(&md5c, (unsigned char *)pPassword, (unsigned int)iPasswordLen);
    MD5Final(bfvec, &md5c);
	memcpy( poRetKey->m_au32Key, bfvec, CHEEZY_SYM_KEY_LEN );
	ImportKey( poRetKey );
	return 0;
}
//---------------------------------------------------------------------------
//! set key used for encryption and decryption
RCODE CheezyCrypt::ImportKey( CheezyKey * poKey )
{
	memcpy( &m_oKey, poKey, sizeof( CheezyKey ) );
	BlowSetKey( &m_oBlowCtx, (unsigned char *)poKey->m_au32Key, CHEEZY_SYM_KEY_LEN );
	m_bIsKeyValid = true;
	return 0;
}


//---------------------------------------------------------------------------
//! Encrypt some data
//! NOTE: iDataLen must be a multiple of 16
RCODE CheezyCrypt::Encrypt( unsigned char * pu8Data, int iDataLen )
{
	if( (false == m_bIsKeyValid) || (iDataLen & 0x0f) )
	{
		return -1;
	}
	BlowEncrypt( &m_oBlowCtx, pu8Data, iDataLen );
	return 0;
}

//---------------------------------------------------------------------------
//! Decrypt some data
//! NOTE: iDataLen must be a multiple of 16
RCODE CheezyCrypt::Decrypt( unsigned char * pu8Data, int iDataLen )
{
	if( (false == m_bIsKeyValid) || (iDataLen & 0x0f) )
	{
		return -1;
	}
	BlowDecrypt( &m_oBlowCtx, pu8Data, iDataLen );
	return 0;
}
//---------------------------------------------------------------------------
//! encrypt known string
RCODE CheezyCrypt::EncryptKnownString( unsigned char * pu8RetData, int iDataLen )
{
	if( (false == m_bIsKeyValid) || (iDataLen != CHEEZY_SYM_KEY_LEN) )
	{
		return -1;
	}
	memcpy( pu8RetData, CHEEZY_KNOWN_TEXT_STR, CHEEZY_SYM_KEY_LEN );
	BlowEncrypt( &m_oBlowCtx, pu8RetData, iDataLen );
	return 0;
}
//---------------------------------------------------------------------------
//! verify data is known text string
RCODE CheezyCrypt::VerifyKnownString( unsigned char * pu8Data, int iDataLen )
{
	if( (false == m_bIsKeyValid) || (iDataLen != CHEEZY_SYM_KEY_LEN) )
	{
		return -1;
	}
	unsigned char au8Buf[ CHEEZY_SYM_KEY_LEN ];
	memcpy( au8Buf, pu8Data, CHEEZY_SYM_KEY_LEN );
	BlowDecrypt( &m_oBlowCtx, au8Buf, CHEEZY_SYM_KEY_LEN );
	return memcmp( au8Buf, CHEEZY_KNOWN_TEXT_STR, CHEEZY_SYM_KEY_LEN );
}

//---------------------------------------------------------------------------
//! fill with random data
//! NOTE: this is not truly random enough for strong encryption uses
void CheezyFillRandom( void * pvData, int iLen )
{
	unsigned long u32Seed = 0;
    #ifdef _MSC_VER
        LARGE_INTEGER u64Timer;
    #else
	    timeval tim;
    #endif

	int i;
	char * pData = (char *)pvData;
	while( 1 )
	{
        #ifdef _MSC_VER
			SYSTEM_INFO sysinfo;
			GetSystemInfo(&sysinfo);
			BOOL bHaveMultipleCPUS = (sysinfo.dwNumberOfProcessors > 1);
			DWORD_PTR affinityMask;
			if(bHaveMultipleCPUS)
			{
				affinityMask = SetThreadAffinityMask(GetCurrentThread(), 1);
			}
            QueryPerformanceCounter(&u64Timer);
			// Restore the true affinity.
			if(bHaveMultipleCPUS)
			{
				(void)SetThreadAffinityMask(GetCurrentThread(), affinityMask);
			}
            u32Seed = u64Timer.LowPart;
        #else
             gettimeofday(&tim, NULL);
             double t1=tim.tv_sec+(tim.tv_usec/1000000.0);
             u32Seed = tim.tv_usec;
        #endif


		unsigned long u32RandCnt = u32Seed & 0x03;
		srand( u32Seed );
		rand();
		if( iLen >= 2 )
		{
			if( u32RandCnt )
			{
				for( i = 0; i < (int)u32RandCnt; i++ )
				{
					* (( short *)pData) = rand();
					while( 0 == * (( short *)pData) )
					{
						* (( short *)pData) = rand();
					}
				}
				iLen -=2;
				pData += 2;
			}
		}
		else
		{
			* pData = (char)rand();

			iLen --;
			pData++;
		}
		if( 0 == iLen )
		{
			return;
		}
		VxSleep( 2 * u32RandCnt );
	}
}

//---------------------------------------------------------------------------
//! Encrypt data with CheezyCrypto
RCODE VxSymEncrypt(		CheezyKey *		poKey,			// Symmetric key must be 16 bytes long
						char *			pDataIn,		// buffer to encrypt
						int				iDataLen,		// data length ( must be multiple of key length )
						char *			pRetBuf )		// if null then encrypted data put in pData
{
	vx_assert( poKey );
	vx_assert( pDataIn );
	vx_assert( iDataLen );
	vx_assert( 0 == (iDataLen & 0x0f) ); // must be multiple of 16 bytes
	char * pDataOut = pRetBuf;
	if( 0 == pDataOut )
	{
		pDataOut = pDataIn;
	}
	CheezyCrypt oCrypto;
	oCrypto.ImportKey( poKey );
	oCrypto.Encrypt( (U8 *)pDataOut, iDataLen );
	return 0;
}
//---------------------------------------------------------------------------
//! Decrypt data with CheezyCrypto
RCODE VxSymDecrypt(		CheezyKey *		poKey,			// Symmetric key must be 16 bytes long
						char *			pDataIn,		// buffer to decrypt
						int				iDataLen,		// data length ( must be multiple of key length )
						char *			pRetBuf )		// if null then encrypted data put in pData
{
	vx_assert( poKey );
	vx_assert( pDataIn );
	vx_assert( iDataLen );
	vx_assert( 0 == (iDataLen & 0x0f) ); // must be multiple of 16 bytes
	char * pDataOut = pRetBuf;
	if( 0 == pDataOut )
	{
		pDataOut = pDataIn;
	}
	CheezyCrypt oCrypto;
	oCrypto.ImportKey( poKey );
	oCrypto.Decrypt( (U8 *)pDataOut, iDataLen );
	return 0;
}
